//
// Created by meowch on 19-11-28.
//

#ifndef LIBDL_CORE_ORDERED_DICT_H
#define LIBDL_CORE_ORDERED_DICT_H

#pragma once

#include "c10/util/Exception.h"
#include "torch/csrc/api/include/torch/ordered_dict.h"
#include "LibDL/utils.h"
#include "./dict_item.h"
#include "boost_any_wrapper.h"

template<typename Key, typename Value>
class Ordered_dict {
public:
    class Item;

    using Iterator = typename std::vector<Item>::iterator;
    using ConstIterator = typename std::vector<Item>::const_iterator;

    using core_type = torch::OrderedDict<Key, Value>;
    core_type core;

    ENABLE_CAST;

    explicit Ordered_dict(const std::string &key_description = "Key") : core(key_description) {};

    Ordered_dict(const Ordered_dict &other) : core(other) {};

    const std::string &key_description() const noexcept {
        return core.key_description();
    };

    Dict_item<Key, Value> front();

    Dict_item<Key, Value> back();

    Value &operator[](const Key &key) {
        return core[key];
    };

    const Value &operator[](const Key &key) const {
        return core[key];
    }

    Value *find(const Key &key) noexcept {
        return core.find(key);
    };

    Value findRef(const Key &key) noexcept {
        return *core.find(key);
    };

    const Value *find(const Key &key) const noexcept(false) {
        return core.find(key);
    };

    Value findRef(const Key &key) const noexcept(false) {
        return *core.find(key);
    };

    bool contains(const Key &key) const noexcept(false) {
        return core.contains(key);
    };

    bool is_empty() const noexcept {
        return core.is_empty();
    };

    // Modifiers
    template<typename K, typename V>
    Value &insert(K &&key, V &&value) {
        return core.insert(key, value);
    };

    Value &insert(Key key, Value &&value) {
        return core.insert(key, value);
    };

    void erase(const Key &key) {
        return core.erase(key);
    };

    void clear() {
        return core.clear();
    };

    // Observers
    std::vector<Dict_item<Key, Value>> items() const noexcept {
        return utils::map_to_wrapper<Dict_item<Key, Value>>(core.items());
    }

    std::vector<Dict_item<Key, Value>> items() noexcept {
        return utils::map_to_wrapper<Dict_item<Key, Value>>(core.items());
    }

    ::std::vector<Key> keys() const {
        return core.keys();
    };

    ::std::vector<Value> values() const {
        return core.values();
    };

private:
    const typename core_type::Item *_front() const {
        return &core.front();
    };

    const typename core_type::Item *_back() {
        return &core.back();
    }
};

//~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~Ordered_dict~~~~~~~~~~~~
template<typename Key, typename Value>
Dict_item<Key, Value> Ordered_dict<Key, Value>::front() {
    return Dict_item<Key, Value>(_front());
}

template<typename Key, typename Value>
Dict_item<Key, Value> Ordered_dict<Key, Value>::back() {
    return Dict_item<Key, Value>(_back());
}

#endif //LIBDL_CORE_ORDERED_DICT_H
